


const fs = require('fs');
const URL = require('url').URL;
const path = require('path');
const http = require("http");

const express = require('express');
const crypto = require('crypto');
const { WebSocketServer } = require("ws");

class HttpServer {
    constructor(API_KEYs_path, public_path) {
        this.app = new express();
        this.app.use(express.json());
        this.app.use(express.urlencoded({ extended: false }));
        this.app.use(express.static(public_path ?? path.join(__dirname, "public")));
        this.API_KEYs = fs.readFileSync(API_KEYs_path ?? path.join(__dirname, './APIKeys.json'));

        this.app.get("/", (req, res) => {
            console.log("\"/\" Called!");
            res.send("Hello World!");
        })

        this.app.get("/randomKeys", (req, res) => {
            console.log("\"/randomKeys\" Called!");
            res.send(crypto.randomBytes(16).toString('hex'));
        });

        setInterval(async function refreshAPI_KEYs() {
            fs.readFile(path.join(__dirname, './APIKeys.json'), (err, data) => {
                this.API_KEYs = JSON.parse(data.toString());
            });
        }, 10 * 60 * 1000);
    }

    async loadAPI_Keys() {
        return new Promise((resolve, reject) => {
            fs.readFile(path.join(__dirname, './APIKeys.json'), (err, data) => {
                if (err) {reject(err)}
                resolve(data);
            });
        })
    }

    get App() {
        return this.app;
    }

    // authorize(str) {
    //     return this.API_KEYs.indexOf(str) !== -1;
    // }
}

class Server extends HttpServer {
    constructor(port, Gemini) {
        super()
        this.port = port;
        this.httpServer = http.createServer(this.App).listen(port);
        this.wss = new WebSocketServer({ noServer: true });
        this.gemini = Gemini;

        const interval = setInterval(() => {
            this.wss.clients.forEach((ws) => {
                if (ws.isAlive === false) return ws.terminate();
                ws.isAlive = false;
                ws.send("Ping");
            });
        }, 30000);

        this.wss.on('close', function close() {
            clearInterval(interval);
        });

        this.httpServer.on('upgrade', (req, socket, head) => {
            socket.on('error', console.error);
            if(req.method !== 'GET') {
                socket.end('HTTP/1.1 405 Method Not Allowed\r\n\r\n');
                return;
            }

            const url = new URL(`ws://localhost:${this.port}/` + req.url);
            if(url.pathname !== '//gemini') {
                socket.end('HTTP/1.1 403 Forbidden\r\n\r\n');
                return;
            }

            if(!url.searchParams.get('auth') || this.API_KEYs.indexOf(url.searchParams.get('auth')) === -1) {
                socket.end('HTTP/1.1 401 Unauthorized\r\n\r\n');
                return;
            }

            this.wss.handleUpgrade(req, socket, head, (ws) => {
                ws.isAlive = true;

                ws.on('error', console.error);
                ws.on('open', () => {
                    console.log("New Ws Connection Established from " + this.getSocketAddress(req));
                    ws.send("HELLO!");
                })
                ws.on('close', () => {
                    ws.removeAllListeners();
                    console.log("Ws Disconnected from " + this.getSocketAddress(socket));
                })

                ws.on('message', (data) => {
                    if(data.toString().toLowerCase() === "pong"){
                        ws.isAlive = true;
                        return;
                    }
                    try { data = JSON.parse(data) }
                    catch(err) { ws.send("Invalid JSON."); return }
                    const { request, type, id, prompt, stamp } = data;

                    if(request !== true) {
                        return;
                    }

                    let ret = {
                        stamp,
                        response: true,
                        result: '',
                        status: 'ERROR',
                    }

                    if (!id) {
                        ret.result = "Invalid ID.";
                        ws.send(JSON.stringify(ret));
                        return;
                    }
                    if(!prompt) {
                        ret.result = "Invalid Prompt";
                        ws.send(JSON.stringify(ret));
                        return;
                    }
                    if(!stamp) {
                        ret.result = "Stamp needed";
                        ws.send(JSON.stringify(ret));
                        return;
                    }

                    switch (type) {
                        case "stream":
                            this.gemini.getReplyStreamWithId(id, prompt)
                                .then(async data => {
                                    let index = 0;
                                    for await (const item of data) {
                                        ret.result = item.text();
                                        ret.status = "OK";
                                        ret.piece = index === 0 ? "start" : "mid";
                                        ws.send(JSON.stringify(ret));
                                        index++;
                                    }
                                    ws.send(JSON.stringify({
                                        ...ret,
                                        result: '',
                                        piece: "finish"
                                    }))
                                })
                                .catch(err => {
                                    console.log(err);
                                    ret.result = err;
                                    ws.send(JSON.stringify(ret));
                                });
                            break;
                        default:
                            this.gemini.getReplyWithId(id, prompt)
                                .then(data => {
                                    ret.result = data;
                                    ret.status = "OK";
                                    ret.piece = "finish";
                                    console.log(data);
                                    ws.send(JSON.stringify(ret));
                                })
                                .catch(err => {
                                    console.log(err);
                                    ret.result = err;
                                    ws.send(JSON.stringify(ret));
                                });
                    }
                })

            });

        });

    }

    getSocketAddress(socket) {
        return `${socket.remoteAddress}:${socket.remotePort}`;
    }
}

module.exports = Server;
